from copy import copy, deepcopy

from common import *

class Action:
    NoAction = 0
    Dream = 1
    Ask = 2

class AI:
    def __init__(self, id):
        self.id = id
        self.epsilon = 0.1
        self.stateActionValues = {}
        self.stateActionValues[(States[0].id, Action.Dream)] = 0.0
        self.stateActionValues[(States[0].id, Action.Ask)] = 0.0
        self.stateActionValues[(States[1].id, Action.Dream)] = 0.5
        self.stateActionValues[(States[1].id, Action.Ask)] = 0.5
        self.reset()
    
    def reset(self):
        self.patterns = []
        self.history = []
        self.activePattern = None
    
    def checkPatternMatch(self, pattern, states, n):
        i = 0
        if n + pattern.length() <= len(states):
            return False
        for i in range(0, pattern.length()):
            if n+i>=len(states):
                return True
            #print i, n
            if not pattern.states[i] == states[n + i]:
                return False
        return False
        
    def checkPatternPredict(self, p, history):
        nHistory = len(history)-1 # First element in history to test
        n = p.length() - 1
        for i in range(0, n):
            if nHistory-i < 0:
                return 0
            if self.checkPatternMatch(p, history, nHistory-i):
                return i+1
        return 0
        
    def predictGetPattern(self, history):
        # Search backwards until a matching pattern begin is found
        nHistory = len(history)-1 # First element in history to test
        if len(history) < 1:
            return (None, 0)
        for p in self.patterns:
            i = self.checkPatternPredict(p, history)
            if i != 0:
                return (p, i)
        if self.activePattern == None:
            return (None, 0)
        # Test active pattern
        i = self.checkPatternPredict(self.activePattern, history)
        if i != 0:
            return (self.activePattern, i)
        return (None, 0)
        
    def predict(self, history):
        (p, i) = self.predictGetPattern(history)
        if p == None:
            return None
        else:
            return p.states[i]
        
    def fitness(self, states):
        correct = 0
        total = 0
        for i in range(1, len(states)):
            if isGoalState(states[i-1]):
                continue
            total = total + 1
            #print states[0:i], self.predict(states[0:i])
            if self.predict(states[0:i]) == states[i]:
                correct = correct + 1
        if total == 0:
            return 0
        return float(correct)/float(total)
    
    def step(self, nextState):
        self.history.append(copy(nextState))
        self.learnPattern()
        self.selectAction()
        if self.action == Action.Dream:
            self.dreamSpecific()
    
    def selectAction(self):
        state = self.history[-1]
        if not isStartState(state):
            self.action = Action.NoAction
            return
        if random.random() < self.epsilon:
            self.action = random.choice([Action.NoAction, Action.Dream, Action.Ask])
            return
        vDream = self.stateActionValues[(state.id, Action.Dream)]
        vAsk = self.stateActionValues[(state.id, Action.Ask)]
        if vDream > vAsk:
            self.action = Action.Dream
        else:
            self.action = Action.Ask
        #print "a: ", self.id, self.history[-1], self.action
    
    def learnFromFitnessAfterAction(self, fitness):
        s = self.history[-1]
        if not isStartState(s):
            return
        id = (s.id, self.action)
        if id in self.stateActionValues.keys():
            old = self.stateActionValues[id]
            self.stateActionValues[id] = old + 0.05*(fitness-old)
        else:
            self.stateActionValues[id] = fitness
        
    
    # Note: Answer invalidates the question.
    def answer(self, question):
        newState = self.predict(question)
        if newState == None:
            return []
        if len(self.history) > 0:
            previousLastHistoryState = self.history[-1]
        else:
            previousLastHistoryState = None
        oldActivePattern = deepcopy(self.activePattern)
        # Append question to history
        for s in question:
            self.addVirtualState(s)
        answer = []
        while newState != None and not isGoalState(newState):
            self.addVirtualState(newState)
            question.append(newState)
            answer.append(deepcopy(newState))
            newState = self.predict(question)
        if not newState == None:
            # Append to history
            self.addVirtualState(newState)
            answer.append(deepcopy(newState))
        if not previousLastHistoryState == None:
            self.history.append(previousLastHistoryState)
            self.activePattern = oldActivePattern
        return answer
    
    def addVirtualState(self, state):
        newState = deepcopy(state)
        newState.virtual = True
        self.history.append(newState)
    
    def dreamSpecific(self):
        states = []
        states.append(random.choice([States[2], States[3]]))
        states.append(random.choice([States[2], States[3]]))
        states.append(States[4])
        self.dream(states)

    def dream(self, states):
        oldState = self.history[-1]
        # Check if we are in a start state
        if not isStartState(oldState):
            return;
        oldActivePattern = deepcopy(self.activePattern)
        #print self.history
        for s in states:
            self.addVirtualState(s)
            self.learnPattern()
        # Revert to pre-dream state
        self.history.append(oldState)
        self.activePattern = oldActivePattern
        #print "o: ", oldActivePattern
        #print self.history
        
    def ask(self, other):
        oldState = self.history[-1]
        # Check if we are in a start state
        if not isStartState(oldState):
            return;
        question = [self.history[-1]]
        #print "q: ", question
        answer = other.answer(question)
        self.dream(answer)
        #print "a: ", answer
    
    def learnPattern(self):
        newState = self.history[-1]
        (predPattern, i) = self.predictGetPattern(self.history[:-1])
        if predPattern != None and not predPattern == self.activePattern and not predPattern.states[i] == newState:
            self.patterns.remove(predPattern)
        if self.activePattern == None:
            self.activePattern = Pattern()
        self.activePattern.states.append(newState)
        # Terminate pattern if goal is reached
        if isGoalState(newState) and self.activePattern.length() > 0:
            # Append to self.patterns if not already there
            self.patterns.append(self.activePattern)
            self.activePattern = None
    
